package pers.qianyu.month_202012.date_20201218;

import org.junit.*;

import java.util.*;
import java.util.stream.*;

/**
 * 经典蓄水池算法
 *
 * @author mizzle rain
 * @date 2020-12-18 20:06
 */
public class ReservoirSampling {
    public List<String> solve(List<String> strs, int K) {
        ArrayList<String> res = new ArrayList<>();
        for (int i = 0; i < strs.size(); i++) {
            String s = strs.get(i);
            if (i < K) {
                res.add(s);
            } else {
                double r = Math.random();
                if (r < ((double) K) / (i + 1)) {
                    int index = (int) (Math.random() * K);
                    res.remove(index);
                    res.add(index, s);
                }
            }
        }
        return res;
    }

    public List<String> solve2(List<String> strs, int K) {
        ArrayList<String> res = new ArrayList<>();
        for (int i = 0; i < strs.size(); i++) {
            String s = strs.get(i);
            if (i < K) {
                res.add(s);
            } else {
                int r = (int) (Math.random() * (i + 1));
                if (r < K) {
                    res.remove(r);
                    res.add(r, s);
                }
            }
        }
        return res;
    }

    @Test
    public void test1() {
        ReservoirSampling rs = new ReservoirSampling();
        List<String> list = Stream.of("a", "b", "c", "d", "e", "f", "g", "h", "i")
                .collect(Collectors.toList());
        for (int i = 0; i < 20; i++) {
            System.out.println(rs.solve(list, 4));
        }
        System.out.println("------------------------------");
        for (int i = 0; i < 20; i++) {
            System.out.println(rs.solve2(list, 4));
        }
    }
}
